###############################################
## This code is for doing Simulation Study III
###############################################
rm(list = ls())


sink("log/log_icew_sim.txt") # start log
begin.time<-Sys.time()


## compile functions 
library(bpNet)
#library(mcmcr)
library(mvtnorm)
library(coda)
set.seed(123456789)

## ------------------------- ##
source("code/simulated30529.R")
source("code/net_plot.R")

## Load the network 
     load('data/SimData/icewsW.RData')
## ------------- Load data   ------------ ##
    load('data/SimData/sim3.RData')
    data1 <- sim3[[1]]$data
    ## true rho (for comparison)
    Rho1 <- sim3[[1]]$Rho
    ## load('data/SimData/sim2cor6.RData')
    ## get the data for simulation II
    rhoZ <- sim3[[2]]


## True parameters
p <- 2
beta <- as.matrix(c(5, 1, 3, -2, 2))

## rho_t
rhoA <- as.matrix(c(0.3, -0.2))
kappa <- 0.3
sigma_n2 <- 0.36

## factor numbers
r <- 2

## error variance
sigma2 <- 1


mcmc <- 15000
burnin <- 5000

## M2: 10 factors with Bayesian shrinkage
out1.2 <- bpNet(data = data1,  
	              W = W,         
	              index = c("id", "time"), 
	              Yname = "y",
	              Xname = c("x1", "x2"),
	              Zname = NULL,   
	              Aname = NULL,   
	              Contextual = c("x1", "x2"),  
	              Contextual.effect = "both",  
	              lagY = FALSE,               
	              force = "both",             
	              r = 10,
	              flasso = 1,             
	              rhoZ = as.matrix(rhoZ),     
	              rho_spec = "ar1",               
                      niter = mcmc,
	              burn = burnin)

## M4: constant rho, 2 factors and no shrinkage
out1.4 <- bpNet(data = data1,   
	              W = W,         
	              index = c("id", "time"), 
	              Yname = "y",
	              Xname = c("x1", "x2"),
	              Zname = NULL,   
	              Aname = NULL,   
	              Contextual = c("x1", "x2"),  
	              Contextual.effect = "both",  
	              lagY = FALSE,                
	              force = "both",              
	              r = 2,
	              flasso = 0,             
	              rhoZ = as.matrix(rhoZ),       
	              rho_spec = "constant",                
	              niter = mcmc,
	              burn = burnin)
 ##------ Making Figure 5 and Figure A8


## ---------------Figure5: plot rho_t
TT <- dim(W)[3]
a <- 25
p1.2 <- rho_plot(out1.2, Rho1, constant = FALSE)
p1.2 <- p1.2 +theme(axis.text.x = element_text( hjust = 1, size=a))
p1.2 <-p1.2 +theme(axis.text.y = element_text( hjust = 1, size=a))
p1.2 <- p1.2 +theme(legend.title = element_text(size=a), legend.text = element_text(size=a), legend.key.size = unit(2, 'cm'))

p1.4 <- rho_plot(out1.4, Rho1, constant = TRUE)
p1.4 <- p1.4 +theme(axis.text.x = element_text( hjust = 1, size=a))
p1.4 <- p1.4 +theme(axis.text.y = element_text( hjust = 1, size=a))
p1.4 <- p1.4 +theme(legend.title = element_text(size=a), legend.text = element_text(size=a), legend.key.size = unit(2, 'cm'))

## Figure 5-(a)
ggsave('plots/sim3/rho2.pdf', p1.2, width = 10, height = 6)

## Figure 5-(b)-i
ggsave('plots/sim3/rho4.pdf', p1.4, width = 10, height = 6)

## Figure 5-(b)-ii
pdf("plots/sim3/rho4hist.pdf",width=6, height=6, pointsize = 15)
rhoC <- out1.4$rho0
CC <- as.mcmc(t(rhoC))
hist(CC, col="gray90", main=" ",  xlab=" ", ylab=" ",  cex.axis=1.4, )
dev.off()


##---- Figure A7: plot omega for factor selection 
p1 <- omegaPlot(out1.2, oxlim = c(-2, 2), plotlength = 5, labelname = "Factor")
ggsave('plots/sim3/factor1.pdf', p1, width = 18, height = 20)

print(Sys.time()-begin.time) 
sink() # close log   












